{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "仿照llama2的GQA代码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import \n",
    "# Copyright (c) Meta Platforms, Inc. and affiliates.\n",
    "# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.\n",
    "\n",
    "import math\n",
    "from dataclasses import dataclass\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "import fairscale.nn.model_parallel.initialize as fs_init\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from fairscale.nn.model_parallel.layers import (\n",
    "    ColumnParallelLinear,\n",
    "    ParallelEmbedding,\n",
    "    RowParallelLinear,\n",
    ")\n",
    "from torch import nn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model setting\n",
    "@dataclass\n",
    "class ModelArgs:\n",
    "    dim: int = 4096\n",
    "    n_layers: int = 32\n",
    "    n_heads: int = 32\n",
    "    n_kv_heads: Optional[int] = None\n",
    "    vocab_size: int = -1  # defined later by tokenizer\n",
    "    multiple_of: int = 256  # make SwiGLU hidden layer size multiple of large power of 2\n",
    "    ffn_dim_multiplier: Optional[float] = None\n",
    "    norm_eps: float = 1e-5\n",
    "\n",
    "    max_batch_size: int = 32\n",
    "    max_seq_len: int = 2048"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 不懂的函数\n",
    "def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):\n",
    "    \"\"\"\n",
    "    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.\n",
    "\n",
    "    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'\n",
    "    and the end index 'end'. The 'theta' parameter scales the frequencies.\n",
    "    The returned tensor contains complex values in complex64 data type.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Dimension of the frequency tensor.\n",
    "        end (int): End index for precomputing frequencies.\n",
    "        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Precomputed frequency tensor with complex exponentials.\n",
    "\n",
    "    \n",
    "        \n",
    "\n",
    "    \"\"\"\n",
    "    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n",
    "    t = torch.arange(end, device=freqs.device)  # type: ignore\n",
    "    freqs = torch.outer(t, freqs).float()  # type: ignore\n",
    "    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64\n",
    "    return freqs_cis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reshape for broadcast\n",
    "\n",
    "\n",
    "def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):\n",
    "    \"\"\"\n",
    "    Reshape frequency tensor for broadcasting it with another tensor.\n",
    "\n",
    "    This function reshapes the frequency tensor to have the same shape as the target tensor 'x'\n",
    "    for the purpose of broadcasting the frequency tensor during element-wise operations.\n",
    "\n",
    "    Args:\n",
    "        freqs_cis (torch.Tensor): Frequency tensor to be reshaped.\n",
    "        x (torch.Tensor): Target tensor for broadcasting compatibility.\n",
    "\n",
    "    Returns:\n",
    "        torch.Tensor: Reshaped frequency tensor.\n",
    "\n",
    "    Raises:\n",
    "        AssertionError: If the frequency tensor doesn't match the expected shape.\n",
    "        AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.\n",
    "    \"\"\"\n",
    "    ndim = x.ndim\n",
    "    assert 0 <= 1 < ndim\n",
    "    assert freqs_cis.shape == (x.shape[1], x.shape[-1])\n",
    "    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n",
    "    return freqs_cis.view(*shape)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RoPE\n",
    "\n",
    "\n",
    "def apply_rotary_emb(\n",
    "    xq: torch.Tensor,\n",
    "    xk: torch.Tensor,\n",
    "    freqs_cis: torch.Tensor,\n",
    ") -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Apply rotary embeddings to input tensors using the given frequency tensor.\n",
    "\n",
    "    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided\n",
    "    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor\n",
    "    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are\n",
    "    returned as real tensors.\n",
    "\n",
    "    Args:\n",
    "        xq (torch.Tensor): Query tensor to apply rotary embeddings.\n",
    "        xk (torch.Tensor): Key tensor to apply rotary embeddings.\n",
    "        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.\n",
    "\n",
    "    Returns:\n",
    "        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.\n",
    "\n",
    "        \n",
    "\n",
    "    \"\"\"\n",
    "    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n",
    "    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n",
    "    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)\n",
    "    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)\n",
    "    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)\n",
    "    return xq_out.type_as(xq), xk_out.type_as(xk)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# repeat_kv\n",
    "def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:\n",
    "    \"\"\"torch.repeat_interleave(x, dim=2, repeats=n_rep)\"\"\"\n",
    "    bs, slen, n_kv_heads, head_dim = x.shape\n",
    "    if n_rep == 1:\n",
    "        return x\n",
    "    return (\n",
    "        x[:, :, :, None, :]\n",
    "        .expand(bs, slen, n_kv_heads, n_rep, head_dim)\n",
    "        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# attention\n",
    "\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    \"\"\"Multi-head attention module.\"\"\"\n",
    "    def __init__(self, args: ModelArgs):\n",
    "        \"\"\"\n",
    "        Initialize the Attention module.\n",
    "\n",
    "        Args:\n",
    "            args (ModelArgs): Model configuration parameters.\n",
    "\n",
    "        Attributes:\n",
    "            n_kv_heads (int): Number of key and value heads.\n",
    "            n_local_heads (int): Number of local query heads.\n",
    "            n_local_kv_heads (int): Number of local key and value heads.\n",
    "            n_rep (int): Number of repetitions for local heads.\n",
    "            head_dim (int): Dimension size of each attention head.\n",
    "            wq (ColumnParallelLinear): Linear transformation for queries.\n",
    "            wk (ColumnParallelLinear): Linear transformation for keys.\n",
    "            wv (ColumnParallelLinear): Linear transformation for values.\n",
    "            wo (RowParallelLinear): Linear transformation for output.\n",
    "            cache_k (torch.Tensor): Cached keys for attention.\n",
    "            cache_v (torch.Tensor): Cached values for attention.\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads\n",
    "        model_parallel_size = fs_init.get_model_parallel_world_size()\n",
    "        self.n_local_heads = args.n_heads // model_parallel_size\n",
    "        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size\n",
    "        self.n_rep = self.n_local_heads // self.n_local_kv_heads\n",
    "        self.head_dim = args.dim // args.n_heads\n",
    "\n",
    "        self.wq = ColumnParallelLinear(\n",
    "            args.dim,\n",
    "            args.n_heads * self.head_dim,\n",
    "            bias=False,\n",
    "            gather_output=False,\n",
    "            init_method=lambda x: x,\n",
    "        )\n",
    "        self.wk = ColumnParallelLinear(\n",
    "            args.dim,\n",
    "            self.n_kv_heads * self.head_dim,\n",
    "            bias=False,\n",
    "            gather_output=False,\n",
    "            init_method=lambda x: x,\n",
    "        )\n",
    "        self.wv = ColumnParallelLinear(\n",
    "            args.dim,\n",
    "            self.n_kv_heads * self.head_dim,\n",
    "            bias=False,\n",
    "            gather_output=False,\n",
    "            init_method=lambda x: x,\n",
    "        )\n",
    "        self.wo = RowParallelLinear(\n",
    "            args.n_heads * self.head_dim,\n",
    "            args.dim,\n",
    "            bias=False,\n",
    "            input_is_parallel=True,\n",
    "            init_method=lambda x: x,\n",
    "        )\n",
    "\n",
    "        self.cache_k = torch.zeros(\n",
    "            (\n",
    "                args.max_batch_size,\n",
    "                args.max_seq_len,\n",
    "                self.n_local_kv_heads,\n",
    "                self.head_dim,\n",
    "            )\n",
    "        ).cuda()\n",
    "        self.cache_v = torch.zeros(\n",
    "            (\n",
    "                args.max_batch_size,\n",
    "                args.max_seq_len,\n",
    "                self.n_local_kv_heads,\n",
    "                self.head_dim,\n",
    "            )\n",
    "        ).cuda()\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        x: torch.Tensor,\n",
    "        start_pos: int,\n",
    "        freqs_cis: torch.Tensor,\n",
    "        mask: Optional[torch.Tensor],\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Forward pass of the attention module.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input tensor.\n",
    "            start_pos (int): Starting position for caching.\n",
    "            freqs_cis (torch.Tensor): Precomputed frequency tensor.\n",
    "            mask (torch.Tensor, optional): Attention mask tensor.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Output tensor after attention.\n",
    "\n",
    "        \"\"\"\n",
    "        bsz, seqlen, _ = x.shape\n",
    "        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)\n",
    "\n",
    "        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)\n",
    "        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)\n",
    "        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)\n",
    "\n",
    "        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)\n",
    "\n",
    "        self.cache_k = self.cache_k.to(xq) # 把cache移动到xq的设备\n",
    "        self.cache_v = self.cache_v.to(xq) # 把cache移动到xq的设备\n",
    "\n",
    "        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk\n",
    "        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv\n",
    "\n",
    "        keys = self.cache_k[:bsz, : start_pos + seqlen]\n",
    "        values = self.cache_v[:bsz, : start_pos + seqlen]\n",
    "\n",
    "        # repeat k/v heads if n_kv_heads < n_heads\n",
    "        keys = repeat_kv(keys, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)\n",
    "        values = repeat_kv(values, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)\n",
    "\n",
    "        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)\n",
    "        keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)\n",
    "        values = values.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)\n",
    "        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)\n",
    "        if mask is not None:\n",
    "            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)\n",
    "        scores = F.softmax(scores.float(), dim=-1).type_as(xq)\n",
    "        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)\n",
    "        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)\n",
    "        return self.wo(output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 前馈feedforward\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim: int,\n",
    "        hidden_dim: int,\n",
    "        multiple_of: int,\n",
    "        ffn_dim_multiplier: Optional[float],\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Initialize the FeedForward module.\n",
    "\n",
    "        Args:\n",
    "            dim (int): Input dimension.\n",
    "            hidden_dim (int): Hidden dimension of the feedforward layer.\n",
    "            multiple_of (int): Value to ensure hidden dimension is a multiple of this value.\n",
    "            ffn_dim_multiplier (float, optional): Custom multiplier for hidden dimension. Defaults to None.\n",
    "\n",
    "        Attributes:\n",
    "            w1 (ColumnParallelLinear): Linear transformation for the first layer.\n",
    "            w2 (RowParallelLinear): Linear transformation for the second layer.\n",
    "            w3 (ColumnParallelLinear): Linear transformation for the third layer.\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        hidden_dim = int(2 * hidden_dim / 3)\n",
    "        # custom dim factor multiplier\n",
    "        if ffn_dim_multiplier is not None:\n",
    "            hidden_dim = int(ffn_dim_multiplier * hidden_dim)\n",
    "        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)\n",
    "\n",
    "        self.w1 = ColumnParallelLinear(\n",
    "            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x\n",
    "        )\n",
    "        self.w2 = RowParallelLinear(\n",
    "            hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x\n",
    "        )\n",
    "        self.w3 = ColumnParallelLinear(\n",
    "            dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.w2(F.silu(self.w1(x)) * self.w3(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RMSNorm\n",
    "\n",
    "class RMSNorm(torch.nn.Module):\n",
    "    def __init__(self, dim: int, eps: float = 1e-6):\n",
    "        \"\"\"\n",
    "        Initialize the RMSNorm normalization layer.\n",
    "\n",
    "        Args:\n",
    "            dim (int): The dimension of the input tensor.\n",
    "            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.\n",
    "\n",
    "        Attributes:\n",
    "            eps (float): A small value added to the denominator for numerical stability.\n",
    "            weight (nn.Parameter): Learnable scaling parameter.\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.weight = nn.Parameter(torch.ones(dim))\n",
    "\n",
    "    def _norm(self, x):\n",
    "        \"\"\"\n",
    "        Apply the RMSNorm normalization to the input tensor.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): The input tensor.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The normalized tensor.\n",
    "\n",
    "        \"\"\"\n",
    "        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass through the RMSNorm layer.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): The input tensor.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The output tensor after applying RMSNorm.\n",
    "\n",
    "        \"\"\"\n",
    "        output = self._norm(x.float()).type_as(x)\n",
    "        return output * self.weight\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transformer block\n",
    "\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, layer_id: int, args: ModelArgs):\n",
    "        \"\"\"\n",
    "        Initialize a TransformerBlock.\n",
    "\n",
    "        Args:\n",
    "            layer_id (int): Identifier for the layer.\n",
    "            args (ModelArgs): Model configuration parameters.\n",
    "\n",
    "        Attributes:\n",
    "            n_heads (int): Number of attention heads.\n",
    "            dim (int): Dimension size of the model.\n",
    "            head_dim (int): Dimension size of each attention head.\n",
    "            attention (Attention): Attention module.\n",
    "            feed_forward (FeedForward): FeedForward module.\n",
    "            layer_id (int): Identifier for the layer.\n",
    "            attention_norm (RMSNorm): Layer normalization for attention output.\n",
    "            ffn_norm (RMSNorm): Layer normalization for feedforward output.\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.n_heads = args.n_heads\n",
    "        self.dim = args.dim\n",
    "        self.head_dim = args.dim // args.n_heads\n",
    "        self.attention = Attention(args)\n",
    "        self.feed_forward = FeedForward(\n",
    "            dim=args.dim,\n",
    "            hidden_dim=4 * args.dim,\n",
    "            multiple_of=args.multiple_of,\n",
    "            ffn_dim_multiplier=args.ffn_dim_multiplier,\n",
    "        )\n",
    "        self.layer_id = layer_id\n",
    "        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        x: torch.Tensor,\n",
    "        start_pos: int,\n",
    "        freqs_cis: torch.Tensor,\n",
    "        mask: Optional[torch.Tensor],\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Perform a forward pass through the TransformerBlock.\n",
    "\n",
    "        Args:\n",
    "            x (torch.Tensor): Input tensor.\n",
    "            start_pos (int): Starting position for attention caching.\n",
    "            freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies.\n",
    "            mask (torch.Tensor, optional): Masking tensor for attention. Defaults to None.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Output tensor after applying attention and feedforward layers.\n",
    "\n",
    "        \"\"\"\n",
    "        h = x + self.attention(\n",
    "            self.attention_norm(x), start_pos, freqs_cis, mask\n",
    "        )\n",
    "        out = h + self.feed_forward(self.ffn_norm(h))\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transformer整体\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, params: ModelArgs):\n",
    "        \"\"\"\n",
    "        Initialize a Transformer model.\n",
    "\n",
    "        Args:\n",
    "            params (ModelArgs): Model configuration parameters.\n",
    "\n",
    "        Attributes:\n",
    "            params (ModelArgs): Model configuration parameters.\n",
    "            vocab_size (int): Vocabulary size.\n",
    "            n_layers (int): Number of layers in the model.\n",
    "            tok_embeddings (ParallelEmbedding): Token embeddings.\n",
    "            layers (torch.nn.ModuleList): List of Transformer blocks.\n",
    "            norm (RMSNorm): Layer normalization for the model output.\n",
    "            output (ColumnParallelLinear): Linear layer for final output.\n",
    "            freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies.\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.params = params\n",
    "        self.vocab_size = params.vocab_size\n",
    "        self.n_layers = params.n_layers\n",
    "\n",
    "        self.tok_embeddings = ParallelEmbedding(\n",
    "            params.vocab_size, params.dim, init_method=lambda x: x\n",
    "        )\n",
    "\n",
    "        self.layers = torch.nn.ModuleList()\n",
    "        for layer_id in range(params.n_layers):\n",
    "            self.layers.append(TransformerBlock(layer_id, params))\n",
    "\n",
    "        self.norm = RMSNorm(params.dim, eps=params.norm_eps)\n",
    "        self.output = ColumnParallelLinear(\n",
    "            params.dim, params.vocab_size, bias=False, init_method=lambda x: x\n",
    "        )\n",
    "\n",
    "        self.freqs_cis = precompute_freqs_cis(\n",
    "            # Note that self.params.max_seq_len is multiplied by 2 because the token limit for the Llama 2 generation of models is 4096. \n",
    "            # Adding this multiplier instead of using 4096 directly allows for dynamism of token lengths while training or fine-tuning.\n",
    "            self.params.dim // self.params.n_heads, self.params.max_seq_len * 2\n",
    "        )\n",
    "\n",
    "    @torch.inference_mode()\n",
    "    def forward(self, tokens: torch.Tensor, start_pos: int):\n",
    "        \"\"\"\n",
    "        Perform a forward pass through the Transformer model.\n",
    "\n",
    "        Args:\n",
    "            tokens (torch.Tensor): Input token indices.\n",
    "            start_pos (int): Starting position for attention caching.\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: Output logits after applying the Transformer model.\n",
    "\n",
    "        \"\"\"\n",
    "        _bsz, seqlen = tokens.shape\n",
    "        h = self.tok_embeddings(tokens)\n",
    "        self.freqs_cis = self.freqs_cis.to(h.device)\n",
    "        freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]\n",
    "\n",
    "        mask = None\n",
    "        if seqlen > 1:\n",
    "            mask = torch.full(\n",
    "                (seqlen, seqlen), float(\"-inf\"), device=tokens.device\n",
    "            )\n",
    "\n",
    "            mask = torch.triu(mask, diagonal=1)\n",
    "\n",
    "            # When performing key-value caching, we compute the attention scores\n",
    "            # only for the new sequence. Thus, the matrix of scores is of size\n",
    "            # (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for\n",
    "            # j > cache_len + i, since row i corresponds to token cache_len + i.\n",
    "            mask = torch.hstack([\n",
    "                torch.zeros((seqlen, start_pos), device=tokens.device),\n",
    "                mask\n",
    "            ]).type_as(h)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            h = layer(h, start_pos, freqs_cis, mask)\n",
    "        h = self.norm(h)\n",
    "        output = self.output(h).float()\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
